from vanna.remote import VannaDefault
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.ZhipuAI.ZhipuAI_Chat import ZhipuAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
from vanna.flask import VannaFlaskApp
from my_openai_chat import My_OpenAI_Chat
from openai import OpenAI
client = OpenAI(
    api_key="Empty",
    base_url="http://localhost:1234/v1"
)


class MyVanna(ChromaDB_VectorStore, My_OpenAI_Chat):
    def __init__(self, client=None, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        My_OpenAI_Chat.__init__(self, client=client, config=config)


vn = MyVanna(client=client, config={"model": "Qwen1.5-7B-Chat-AWQ-fp16", "temperature": 0.5, "max_tokens": 1000})
vn.connect_to_postgres(host='10.10.50.89', dbname='smartpole', user='smartpole', password='Telchina@2021', port='57287')

df_information_schema = vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS")

# This will break up the information schema into bite-sized chunks that can be referenced by the LLM
plan = vn.get_training_plan_generic(df_information_schema)
vn.train(plan=plan)
# vn.train(
#     question="都有哪些视频点信息？",
#     sql="SELECT * FROM device.device_camera"
# )
# DDL statements are powerful because they specify table names, colume names, types, and potentially relationships
# vn.train(ddl="""
# CREATE TABLE "device"."device_broadcast" (
#   "id" varchar(64) COLLATE "pg_catalog"."default" NOT NULL,
#   "lamp_post_id" varchar(64) COLLATE "pg_catalog"."default",
#   "name" varchar(100) COLLATE "pg_catalog"."default",
#   "code" varchar(100) COLLATE "pg_catalog"."default",
#   "broadcast_out_volume" int4,
#   "device_model_id" varchar(64) COLLATE "pg_catalog"."default",
#   "company_id" varchar(64) COLLATE "pg_catalog"."default",
#   "online_state" varchar(50) COLLATE "pg_catalog"."default",
#   "register_state" varchar(50) COLLATE "pg_catalog"."default",
#   "insert_time" timestamp(6),
#   "update_time" timestamp(6),
#   "comment" varchar(500) COLLATE "pg_catalog"."default",
#   "organ_id" varchar(64) COLLATE "pg_catalog"."default",
#   "external_device_id" varchar COLLATE "pg_catalog"."default",
#   PRIMARY KEY ("id")
# )
# ;
#
# ALTER TABLE "device"."device_broadcast"
#   OWNER TO "smartpole";
#
# COMMENT ON COLUMN "device"."device_broadcast"."id" IS '主键';
#
# COMMENT ON COLUMN "device"."device_broadcast"."lamp_post_id" IS '关联灯杆';
#
# COMMENT ON COLUMN "device"."device_broadcast"."name" IS '名称';
#
# COMMENT ON COLUMN "device"."device_broadcast"."code" IS '编号';
#
# COMMENT ON COLUMN "device"."device_broadcast"."broadcast_out_volume" IS '广播输出音量';
#
# COMMENT ON COLUMN "device"."device_broadcast"."device_model_id" IS '设备型号';
#
# COMMENT ON COLUMN "device"."device_broadcast"."company_id" IS '生产厂商';
#
# COMMENT ON COLUMN "device"."device_broadcast"."online_state" IS '设备状态';
#
# COMMENT ON COLUMN "device"."device_broadcast"."register_state" IS '注册状态';
#
# COMMENT ON COLUMN "device"."device_broadcast"."insert_time" IS '插入时间';
#
# COMMENT ON COLUMN "device"."device_broadcast"."update_time" IS '更新时间';
#
# COMMENT ON COLUMN "device"."device_broadcast"."comment" IS '备注';
#
# COMMENT ON COLUMN "device"."device_broadcast"."organ_id" IS '所属单位';
#
# COMMENT ON COLUMN "device"."device_broadcast"."external_device_id" IS '外部设备 Id';
#
# COMMENT ON TABLE "device"."device_broadcast" IS '公共广播管理';
# """)

# Sometimes you may want to add documentation about your business terminology or definitions.
#vn.train(documentation="查询")

# You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
#vn.train(sql="SELECT * FROM device.device_broadcast")

# At any time you can inspect what training data the package is able to reference
#training_data = vn.get_training_data()
#training_data

# You can remove training data if there's obsolete/incorrect information.
#vn.remove_training_data(id='1-ddl')

## Asking the AI
#vn.ask(question="")


app = VannaFlaskApp(vn)
app.run()
